import os
import gymnasium as gym
from train import REINFORCE_Agent
from config import (GAME_NAME, HIDDEN_DIM, DEVICE, LEARNING_RATE, DISCOUNT_FACTOR)

env = gym.make(GAME_NAME, render_mode="human", max_episode_steps=1500)
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
action_low = env.action_space.low
action_high = env.action_space.high

# 创建智能体
agent = REINFORCE_Agent(
    state_dim=state_dim,
    hidden_dim=HIDDEN_DIM,
    action_dim=action_dim,
    device=DEVICE,
    learning_rate=LEARNING_RATE,
    discount_factor=DISCOUNT_FACTOR,
    action_low=action_low,
    action_high=action_high,
)
# 加载已有模型
results_dir = os.path.join("results", GAME_NAME)
os.makedirs(results_dir, exist_ok=True)
model_path = os.path.join(results_dir, "reinforce_model.pth")
agent.load_model(model_path)

for ep in range(5):
    episode_return = 0.0
    done = False
    count = 0
    state, _ = env.reset()
    while not done:
        action, _ = agent.take_action(state)
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        state = next_state
        episode_return += reward
        count += 1
    print(f"Episode {ep + 1}: Return = {episode_return:.2f}")
    print(f"count:{count}")
env.close()

